# tasks.py
from django.utils import timezone
from django.core.files.storage import FileSystemStorage
from django.core.cache import cache
from django.db.models import F
from django.core.files.base import ContentFile

# from __future__ import absolute_import, unicode_literals
from celery import shared_task
from openai import OpenAI
from PIL import Image

import json
import os
import requests
import time
import datetime

# from datetime import datetime
from urllib.parse import urlparse, parse_qs

from ...models.users import Users
from ...models.api_keys import ApiKeys
from ...models.dall_jobs import DallJob
from ...models.powerlogs import PowerLogs

from ...utils.users import rand_num_file

from CeleryProject import app


from asgiref.sync import async_to_sync
from channels.layers import get_channel_layer
def some_view(jobid, userid, res):

    channel_layer = get_channel_layer()

    # 发送信息到websocket
    if res["code"] == 1:
        # 等于1 代表任务里面有失败 更新数据库
        # 更新任务进度
        DallJob.objects.filter(id=jobid).update(
                                            progress = -1,
                                            errmsg = res["mes"]
                                        )
        
        async_to_sync(channel_layer.group_send)(f'user_{userid}', {
            'type': 'my_event',
            'message': "FAIL"
        })
    else:
        mes = res["mes"]
        async_to_sync(channel_layer.group_send)(f'user_{userid}', {
            'type': 'my_event',
            'message': mes
        })


@shared_task()
def testtest(x, y):
    import time
    time.sleep(10)
    return x+y


@shared_task
def dalltasks(jobid, userid, prompt, quality, size, style, dallpower):
    # 获取用户的当前剩余的算力
    user_instance = Users.objects.get(id=userid)
    if user_instance.power < dallpower:
        res = {
            "code": 1,
            "mes": "算力不足, 需要：%s, 剩余：%s" % (dallpower, user_instance.power)
        }
        some_view(jobid, userid, res)
        return res
    
    # 获取绘画api key
    apikey_instance = ApiKeys.objects.filter(
                        platform = "OpenAI",
                        type = "img",
                        enabled = True
                    ).first()
    api_key = apikey_instance.value
    if not apikey_instance:
        # 没有可用的绘画api key
        res = {
            "code": 1,
            "mes": "没有可用的IMG api密钥"
        }
        some_view(jobid, userid, res)
        return res

    # 获取api key的 url
    # 获取api url
    api_url = apikey_instance.api_url
    if not api_url:
        api_url = f"https://api.openai.com/v1"

    # 请求openai
    res = stream_openai_response(api_key, api_url, prompt, size, quality, style)
    # time.sleep(10)
    # res = {
    #     "code": 0,
    #     "mes": "https://gips0.baidu.com/it/u=3602773692,1512483864&fm=3028&app=3028&f=JPEG&fmt=auto?w=960&h=1280"
    # }

    if res["code"] == 0:
         org_url = res["mes"]
    else:
        some_view(jobid, userid, res)
        return res

    # 更新api key的最后使用时间
    now = timezone.now()
    apikey_instance.last_used_at = now
    apikey_instance.save()

    # 更新任务进度
    DallJob.objects.filter(id=jobid).update(
                                        progress = 100,
                                        orgurl = org_url,
                                        prompt = prompt
                                    )

    # res = {
    #     "code": 0,
    #     "mes": "FINISH"
    # }
    # # 发送成功的消息
    # some_view(jobid, userid, res)
    
    ###########################################
    # 下载图片，因为openai生成的图片 链接 1小时内有效

    # 先获取图片格式
    # 解析url
    parsed_url = urlparse(org_url)
    
    # 获取查询参数字符串
    query_params = parsed_url.query
    
    # 解析查询参数
    params = parse_qs(query_params)
    # uploaded_file_ext = params['rsct'][0].split("/")[1]
    uploaded_file_ext = ".png"


    # 下载图片二进制
    try:
        generated_image = requests.get(org_url).content

    except:
        res = {
            "code": 1,
            "mes": "下载图片失败！图片已经过期！"
        }
        some_view(jobid, userid, res)
        return res

    time_path = "dallimages/" + datetime.datetime.now().strftime("%Y/%m/%d")

    # 创建保存文件的实例
    fs = FileSystemStorage()

    # 创建保存文件的基础路径
    os.makedirs(os.path.join(fs.location, time_path), exist_ok=True)

    # 随机生成16位数字文件名称
    rand_file_name = rand_num_file(uploaded_file_ext)

    # 获取保存文件的路径
    file_path = os.path.join(time_path, rand_file_name)

    # 保存文件
    try:
        filename = fs.save(file_path, ContentFile(generated_image))
        visit_url = fs.url(filename)
    except:
        res = {
            "code": 1,
            "mes": "保存图片失败！"
        }
        some_view(jobid, userid, res)
        return res
    
    ########################################################################################################################
    # 获取高清图片的缩略图
    # 缩略图会保存下来，记得删除的时候 一并删除

    # 打开高清图片
    # 先得到
    try:
        myimage = Image.open(visit_url[1:])

        # 设置缩略图的大小
        width, height = myimage.size

        # 假设我们想要的缩略图大小是300x200
        thumb_size = (300, 200)

        # 确保不失真缩放图片
        thumb_size = (width if width < thumb_size[0] else thumb_size[0],
                    height if height < thumb_size[1] else thumb_size[1])
        
        # 生成缩略图
        myimage.thumbnail(thumb_size)

        # 保存缩略图到同一目录，文件名后加_thumb
        # 先获取缩略图文件名称
        thumb_name, thumb_ext = os.path.splitext(rand_file_name)
        thumb_filename = f"{thumb_name}_thumb{thumb_ext}"

        # 获取保存文件的相对路径
        thumburl_path = os.path.join("/media/", time_path, thumb_filename)
        file_thumb_path = os.path.join("media/", time_path, thumb_filename)

        myimage.save(file_thumb_path, 'JPEG')
    except:
        pass


    ########################################################################################################################

    # 更新job的信息 img url 
    # /media/dallimages/2024/12/16/4992544180653704.png
    # /media/dallimages/2024/12/16/4992544180653704_thumb.png
    DallJob.objects.filter(id=jobid).update(imgurl = visit_url, thumburl = thumburl_path)


    # 等这个任务完全成功了之后 再发送完成的信息
    res = {
        "code": 0,
        "mes": "FINISH"
    }
    # 发送成功的消息
    some_view(jobid, userid, res)

    # 更新用户算力
    # 扣减算力 放到前面所有的操作都 弄完后。这样避免了 任务不成功 然后出现扣减算力的情况
    Users.objects.filter(id=userid).update(power=F('power') - dallpower)

    user_update_obj = Users.objects.get(id=userid)

    # 更新用户信息缓存
    user_cache_mes = cache.get(f'user_info_{userid}')
    if user_cache_mes:
        user_cache_mes['power'] = user_update_obj.power
        cache.set(f'user_info_{userid}', user_cache_mes, timeout=3600)

    
    # 记录算力变化日志
    # 记录消费日志
    if len(prompt) > 10:
        remark_prompt = prompt[:10] + "..."
    else:
        remark_prompt = prompt


    remark_str = "绘画提示词: {}".format(remark_prompt)
    powerlogs_obj = PowerLogs(
        user_id = userid,
        username = user_update_obj.username,
        type = 2,
        amount = dallpower,
        balance = user_update_obj.power,
        model = "dall-e-3",
        remark = remark_str,
        mark = 0
    )

    try:
        powerlogs_obj.save()
    except Exception as e:
        print("记录消费日志失败！")


    content = "```\n%s\n```\n下面是我为你创作的图片：\n\n![](%s)\n" % (prompt, visit_url)
    res = {
        "code": 0,
        "mes": content
    }
    return res


def stream_openai_response(api_key, api_url, u_prompt, u_size, u_quality, u_style):
        client = OpenAI(
            api_key= api_key, # 如果您没有配置环境变量，请在此处用您的API Key进行替换
            base_url= api_url,  # 填写DashScope服务的base_url
        )

        try:
            response = client.images.generate(
                model="dall-e-3",
                prompt=u_prompt,
                size=u_size,
                quality=u_quality,
                n=1,
                style=u_style
            )
            res = {
                "code": 0,
                "mes": response.data[0].url
            }
            return res
        except Exception as e:
            erro_mes_bits = e.response.content
            erro_mes_dict = json.loads(erro_mes_bits)
            erro_mes = erro_mes_dict["error"]["message"]
            res = {
                "code": 1,
                "mes": erro_mes
            }
            return res


